#include "FindTransform.h"
#include <algorithm>
#include <opencv2\nonfree\features2d.hpp>
#include "support.h"
#include <QDebug>

FindTransform::FindTransform(void) : Filter(INPUTS_COUNT, 1)
{
	extractor = new cv::SurfDescriptorExtractor();
//	matcher = new cv::BFMatcher(cv::NORM_L2);
	matcher = new cv::FlannBasedMatcher();
}


FindTransform::~FindTransform(void)
{
	delete extractor;
	delete matcher;
}

void FindTransform::Process()
{
	/*
		input0 - last keypoints
		input1 - current keypoints
		input2 - last image
		input3 - current image
		input4 - last points3d
		input5 - current points3d

		output0 - transformation matrix
	*/

	std::vector<cv::KeyPoint> *lastKeypoints = new std::vector<cv::KeyPoint>();
	std::vector<cv::KeyPoint> *currentKeypoints = new std::vector<cv::KeyPoint>();

	for (int i = 0; i < inputs[0]->rows; ++i)
	{
		float *ptr = inputs[0]->ptr<float>(i);
		cv::KeyPoint kp;
		kp.pt.x = ptr[0];
		kp.pt.y = ptr[1];
		kp.size = ptr[2];
		kp.angle = ptr[3];
		kp.response = ptr[4];
		kp.octave = *reinterpret_cast<int*>(ptr + 5);
		kp.class_id = *reinterpret_cast<int*>(ptr + 6);

		lastKeypoints->push_back(kp);
	}

	for (int i = 0; i < inputs[1]->rows; ++i)
	{
		float *ptr = inputs[1]->ptr<float>(i);
		cv::KeyPoint kp;
		kp.pt.x = ptr[0];
		kp.pt.y = ptr[1];
		kp.size = ptr[2];
		kp.angle = ptr[3];
		kp.response = ptr[4];
		kp.octave = *reinterpret_cast<int*>(ptr + 5);
		kp.class_id = *reinterpret_cast<int*>(ptr + 6);

		currentKeypoints->push_back(kp);
	}

	//extract descriptors
	
	cv::Mat lastDescriptors;
	cv::Mat currentDescriptors;

	extractor->compute(*inputs[2], *lastKeypoints, lastDescriptors);
	extractor->compute(*inputs[3], *currentKeypoints, currentDescriptors);

	//match keypoints
	std::vector<cv::DMatch> matches;
	matcher->match(currentDescriptors, lastDescriptors, matches);

	delete lastKeypoints;
	delete currentKeypoints;

	//remove outliers
	float lavgz = 0.0f;
	float cavgz = 0.0f;

	int ptscount = 0;

	for (auto it = matches.begin(); it != matches.end(); ++it)
	{
		float lz = inputs[4]->at<float>(it->trainIdx, 2);
		float cz = inputs[5]->at<float>(it->queryIdx, 2);

		if (lz >= 100 && lz <= 1000)
		{
			lavgz += lz;
			cavgz += cz;
			ptscount ++;
		}
	}

	lavgz /= ptscount ;
	cavgz /= ptscount ;

	int removed = 0;
	for (int i = 0; i < matches.size(); ++i)
	{
		float lz = inputs[4]->at<float>(matches[i].trainIdx, 2);

		if (abs(lz - lavgz) > 300)
		{
			matches.erase(matches.begin() + i);
			removed++;
			--i;
			continue;
		}

		float cz = inputs[5]->at<float>(matches[i].queryIdx, 2);
		if (abs(cz - cavgz) > 300)
		{
			matches.erase(matches.begin() + i);
			removed++;
			i--;
		}
	}

//	qDebug() << "Removed: " << removed;

	//find centroids
	cv::Point3f centroidA = cv::Point3f();
	cv::Point3f centroidB = cv::Point3f();

	//select first 10
	std::sort(matches.begin(), matches.end());

	int count = MIN(5, matches.size());
	matches = std::vector<cv::DMatch>(matches.begin(), matches.begin() + count);

	std::vector<cv::Point3f> lastPoints;
	std::vector<cv::Point3f> currentPoints;

	for (auto it = matches.begin(); it != matches.end(); ++it)
	{
		cv::Point3f lastPoint;
		cv::Point3f currentPoint;

		float *lp = inputs[4]->ptr<float>(it->trainIdx);
		float *cp = inputs[5]->ptr<float>(it->queryIdx);

		lastPoint.x = lp[0];
		lastPoint.y = lp[1];
		lastPoint.z = lp[2];

		currentPoint.x = cp[0];
		currentPoint.y = cp[1];
		currentPoint.z = cp[2];

		centroidA += lastPoint;
		centroidB += currentPoint;

		lastPoints.push_back(lastPoint);
		currentPoints.push_back(currentPoint);
	}

	centroidA *= 1/(float)count;
	centroidB *= 1/(float)count;

	cv::Mat H = cv::Mat(3,3,CV_32FC1);

	for (int i = 0; i < lastPoints.size(); ++i)
	{
		cv::Point3f lastPoint = lastPoints[i];
		cv::Point3f currentPoint = currentPoints[i];

		cv::Mat lp = cv::Mat(lastPoint - centroidA);
		cv::Mat cp = cv::Mat(currentPoint - centroidB).t();

		cv::Mat product = lp * cp;
		H += product;
	}

	cv::SVD svd(H);

	cv::Mat V = svd.vt.t();
	cv::Mat Ut = svd.u.t();

	cv::Mat R = V * Ut;

	cv::Mat Ca = cv::Mat::eye(4,4,CV_32FC1);
	cv::Mat Rnew = cv::Mat::eye(4,4,CV_32FC1);
	cv::Mat Cb = cv::Mat::eye(4,4,CV_32FC1);

	cv::Mat centroidAMat = cv::Mat(centroidA);
	cv::Mat centroidBMat = cv::Mat(centroidB);

	Ca(cv::Range(0, 3), cv::Range(3,4)) = -centroidAMat;
	Cb(cv::Range(0, 3), cv::Range(3,4)) = (-(-centroidBMat));
	R.copyTo(Rnew(cv::Rect(0, 0, 3, 3)));

	outputs[0] = Cb * Rnew * Ca;

//	printMat(outputs[0]);
}
